import os
import json
from collections import Counter

os.chdir('..')
cur = os.getcwd()
print('当前数据处理默认工作目录：', cur)


class TransferData():
    def __init__(self):
        self.label_dict = json.load(open(os.path.join(cur, 'data/data/labels.json'), encoding='utf-8'))
        self.seq_tag_dict = json.load(open(os.path.join(cur, 'data/data/tag2id.json'), encoding='utf-8'))
        self.train_filepath = os.path.join(cur, 'data/data/train.txt')
        self.origin_path = os.path.join(cur, 'data/data_origin')

    def transfer(self):
        with open(self.train_filepath, 'w', encoding='utf-8') as fr:
            #     print('jejej')
            #     print(self.origin_path)
            #     print(list(os.walk(self.origin_path)))
            for root, dirs, files in os.walk(self.origin_path):
                for file in files:
                    filepath = os.path.join(root, file)
                    if 'original' not in filepath:
                        continue
                    label_filepath = filepath.replace('.txtoriginal', '')
                    # print(filepath, '\t\t', label_filepath)
                    res_dict = self.read_label_text(label_filepath)
                    with open(filepath, 'r', encoding='utf-8') as f:
                        content = f.read().strip()
                        for indx, char in enumerate(content):
                            char_label = res_dict.get(indx, 'O')
                            fr.write(char + '\t' + char_label + '\n')

    def read_label_text(self, label_filepath):
        res_dict = {}
        for line in open(label_filepath, 'r', encoding='utf-8'):
            res = line.strip().split('\t')
            start = int(res[1])
            end = int(res[2])
            label = res[3]
            label_tag = self.label_dict.get(label)
            for i in range(start, end + 1):
                if i == start:
                    tag = 'B-' + label_tag
                else:
                    tag = 'I-' + label_tag
                res_dict[i] = tag
        return res_dict


# D:\HeiMa\pythonProject3\NER\data\data\labels.json    open(os.path.join(cur, 'data/labels.json')
if __name__ == '__main__':
    t = TransferData()
    # print(t.label_dict)
    t.transfer()
